import mxnet as mx
import argparse
import importlib
import train_model


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--network', type=str, default='vgg',
                        help='the cnn to use')
    parser.add_argument('--gpus', type=str,
                        help='the gpus will be used, e.g "0,1,2,3"')
    parser.add_argument('--batch-size', type=int, default=128,
                        help='the batch size')
    parser.add_argument('--lr', type=float, default=.1,
                        help='the initial learning rate')
    parser.add_argument('--model-prefix', type=str,
                        help='the prefix of the model to load/save')
    parser.add_argument('--save-model-prefix', type=str,
                        help='the prefix of the model to save')
    parser.add_argument('--num-epochs', type=int, default=10,
                        help='the number of training epochs')
    parser.add_argument('--load-epoch', type=int,
                        help="load the model on an epoch using the model-prefix")
    parser.add_argument('--lr-factor', type=float, default=1,
                        help='times the lr with a factor for every lr-factor-epoch epoch')
    parser.add_argument('--lr-factor-epoch', type=float, default=1,
                        help='the number of epoch to factor the lr, could be .5')
    parser.add_argument('--kv-store', type=str, default='local',
                        help='the kvstore type')
    parser.add_argument('--num-examples', type=int, default=10000,
                        help='number of training examples')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    if args.save_model_prefix is None:
        args.save_model_prefix = "DP_UP_" + args.network

    net = importlib.import_module('symbol_' + args.network).get_symbol(num_classes=2)

    train_dataiter = mx.io.ImageRecordIter(
        path_imgrec="rec/beijing-0-1-train.rec",
        mean_img="mean.bin",
        data_shape=(3, 230, 230),
        batch_size=args.batch_size,
        rand_crop=True,
        rand_mirror=True,
        shuffle=True,
        preprocess_threads=4,
        prefetch_buffer=1000)

    test_dataiter = mx.io.ImageRecordIter(
        path_imgrec="rec/beijing-0-1-test.rec",
        mean_img="mean.bin",
        data_shape=(3, 230, 230),
        batch_size=args.batch_size,
        rand_crop=False,
        rand_mirror=False,
        shuffle=False,
        preprocess_threads=4,
        prefetch_buffer=1000)

    # batch = train_dataiter.next().data[0].asnumpy()
    # from img.vis import SavePicFromNumpy
    # for i in range(3):
    #    SavePicFromNumpy(batch[i], 'img/%d.jpg' % i)

    train_model.fit(args, net, train_dataiter, test_dataiter)
